package cn.doitedu.demo7;

import cn.doitedu.beans.UserAction;
import groovy.lang.GroovyClassLoader;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.HashMap;
import java.util.Map;

@Slf4j
public class CoreProcessFunction_OperatorState extends KeyedBroadcastProcessFunction<Long, UserAction, RuleMeta, String> implements CheckpointedFunction {

    HashMap<String, RuleCalculator> ruleCalculatorPool;
    //MapState<String, RuleCalculator> ruleCalculatorPool;

    // 最近2分钟的用户行为明细缓存
    ListState<UserAction> recentActions;

    // 用户存放已处理过数据的运算机的规则id
    MapState<String, String> noNewCalculatorIds;

    GroovyClassLoader groovyClassLoader;

    ListState<RuleCalculator> calculatorOperatorState;

    @Override
    public void open(Configuration parameters) throws Exception {
        ruleCalculatorPool = new HashMap<>();
        //ruleCalculatorPool = getRuntimeContext().getMapState(new MapStateDescriptor<String, RuleCalculator>("calculator-pool", String.class, RuleCalculator.class));


        ListStateDescriptor<UserAction> desc = new ListStateDescriptor<>("recent_actions", UserAction.class);
        StateTtlConfig ttlConfig = StateTtlConfig.newBuilder(Time.minutes(2))
                .neverReturnExpired()
                .build();
        desc.enableTimeToLive(ttlConfig);

        recentActions = getRuntimeContext().getListState(desc);


        this.noNewCalculatorIds = getRuntimeContext().getMapState(new MapStateDescriptor<String, String>("no_new_calculator_ids", String.class, String.class));


        this.groovyClassLoader = new GroovyClassLoader();

    }

    @Override
    public void processElement(UserAction currentAction, KeyedBroadcastProcessFunction<Long, UserAction, RuleMeta, String>.ReadOnlyContext ctx, Collector<String> out) throws Exception {


        // 先把本次的用户行为放入缓存状态
        recentActions.add(currentAction);


        // 遍历运算机池
        for (Map.Entry<String, RuleCalculator> entry : ruleCalculatorPool.entrySet()) {

            // 取到一个规则运算机对象
            String ruleId = entry.getKey();
            RuleCalculator ruleCalculator = entry.getValue();

            // 判断该运算机是否是一个新上线的运算机
            if (!noNewCalculatorIds.contains(ruleId)) {
                // 如果是，则要把缓存状态中的最近2分钟历史数据交给这个运算机处理
                for (UserAction action : recentActions.get()) {
                    ruleCalculator.calc(action, out);
                }

                // 然后将这个新运算加入 “老运算机id集合”
                noNewCalculatorIds.put(ruleId, "");

            } else {
                // 如果否,那就把本次收到的这一条行为交给它处理
                ruleCalculator.calc(currentAction, out);
            }
        }

    }

    @Override
    public void processBroadcastElement(RuleMeta ruleMeta, KeyedBroadcastProcessFunction<Long, UserAction, RuleMeta, String>.Context ctx, Collector<String> out) throws Exception {
        String op = ruleMeta.getOp();

        // 如果是元数据删除操作，或 元数据更新操作且更新后的规则状态为0 ，都表示要下线规则
        if (op.equals("d") || (op.equals("u") && ruleMeta.getRule_status() == 0)) {
            ruleCalculatorPool.remove(ruleMeta.getRule_id());

            log.warn("下线了一个规则运算机,rule_id:{}", ruleMeta.getRule_id());

        }
        // 如果是初始读或新增了一条规则元数据，且规则的状态为1，表示要上线规则
        else if ((op.equals("r") || op.equals("c") || op.equals("u")) && ruleMeta.getRule_status() == 1) {
            // 根据规则元数据，构造运算机对象
            String ruleModelId = ruleMeta.getRule_model_id();


            // 从规则元数据中，取出规则对应的模型运算机groovy代码
            String modelCode = ruleMeta.getModel_code();
            Class<?> aClass = groovyClassLoader.parseClass(modelCode);

            RuleCalculator ruleCalculator = (RuleCalculator) aClass.newInstance();


            // 并利用元数据中的 规则参数，初始化运算机，得到具体规则运算机对象
            ruleCalculator.init(getRuntimeContext(), ruleMeta);

            // 将初始化好的运算机对象，放入运算机池
            ruleCalculatorPool.put(ruleMeta.getRule_id(), ruleCalculator);

            log.info("上线了一个规则运算机,rule_id:{}", ruleMeta.getRule_id());


        } else {
            // TODO :  NOTHING TO DO
            log.info("收到一个规则元数据，但我什么也没做,rule_id:{}", ruleMeta.getRule_id());
        }
    }


    /**
     * checkpoint做快照时，会先调用这个方法
     */
    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {

        for (Map.Entry<String, RuleCalculator> entry : ruleCalculatorPool.entrySet()) {

            RuleCalculator ruleCalculator = entry.getValue();
            // 把它装入operatorState
            calculatorOperatorState.add(ruleCalculator);
        }


    }

    /**
     * failover恢复，或者是从某个检查点人工启动
     */
    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {

        calculatorOperatorState = context.getOperatorStateStore().getUnionListState(new ListStateDescriptor<RuleCalculator>("calculator_state", RuleCalculator.class));

        for (RuleCalculator ruleCalculator : calculatorOperatorState.get()) {

            ruleCalculatorPool.put(ruleCalculator.ruleId(),ruleCalculator);

        }


    }

}
